
import torch
from torch.autograd import Variable
import torchvision
from torchvision import datasets, transforms, models
import numpy
from tqdm import tqdm
import os
import time

resnet_50_full = models.resnet50(pretrained=True)

resnet_50_full = resnet_50_full.cuda()

resnet_layer = torch.nn.Sequential(*list(resnet_50_full.children())[:-2])
# print(resnet_layer)

input_data = torch.rand(64,3,224,224)
input_data = input_data.cuda()

output = resnet_layer(input_data)

print(output.shape)


output = resnet_50_full(input_data)

print(output.shape)